# coding: utf-8
import asyncio
import shutil
from datetime import datetime

from httpx import AsyncClient
from loguru import logger

from common import signalBus, TaskStatus, Task, TaskMessage, COVER_DIR
from .m3u8_tools import redirected_resolution_url, m3u8_to_absolute_uris, DecodeByte, merge_ts_to_mp4
from .runner_base import DownRunnableBase
from .tools import speed_to_str, get_video_cover


class M3u8Runner(DownRunnableBase):
    def __init__(
            self,
            task: Task
    ):
        super().__init__(task)
        self.timeout = 10
        self.isStop = False

        self.loop: asyncio.AbstractEventLoop = None  # 事件循环
        self.start_time = datetime.now()  # 开始时间
        self.download_num = 0  # 已下载数据块数
        self.total_num = 0  # 总数据数
        self.download_size = 0  # 已下载数据大小 单位：字节
        self.ts_urls = []  # ts文件下载地址
        self.decrypt_key = {}  # 解密key
        self.error_string = ""  # 错误信息
        self.__mainTask: asyncio.Task = None  # 主任务

        parent_dir = self.save_file.parent
        self.temp_dir = parent_dir / f'{self.save_file.stem}_temp'
        self.ffmpeg_file = self.temp_dir / 'ffmpeg.txt'
        self.taskMsg = TaskMessage(
            id=self.tid,
        )

    async def fetch(
            self,
            ts_url: str,
            ts_name: str,
            key: dict,
            client: AsyncClient,
            *,
            retries: int = 0,

    ):
        try:
            if not (self.temp_dir / ts_name).exists():
                iv = key.get('iv', ts_name.split('.')[0].zfill(16))
                response = await client.get(ts_url)
                response.raise_for_status()
                content = response.content
                if key.get("uri") and key.get("method") and key.get("key"):
                    decoded_content = DecodeByte.do_decode(key["key"], iv, content, method=key.get("method", "AES-128"))
                else:
                    decoded_content = content
                # 保存ts文件
                with open(self.temp_dir / ts_name, 'wb') as f:
                    f.write(decoded_content)
                    f.flush()

                # 保存ffmpeg文件
                with open(self.ffmpeg_file, 'a', encoding='utf-8') as f:
                    f.write(f"file '{ts_name}'\n")
                    f.flush()
            else:
                decoded_content = ''
            self.download_num += 1
            self.download_size += len(decoded_content)
            # 发送进度信号
            progress = (self.download_num / self.total_num) * 100
            self.taskMsg.progress = progress
            self.taskMsg.downloaded_msg = f'{self.download_num} / {self.total_num}'
            self.taskMsg.message = f"下载进度：{progress:.2f}%"

            elapsed_time = (datetime.now() - self.start_time).total_seconds()  # 计算已用时间
            if elapsed_time > 0:  # 避免除以零的错误
                speed = self.download_size / elapsed_time  # 当前速度（字节/秒）
                self.taskMsg.net_speed = f'{speed_to_str(speed)}/S'
            signalBus.taskMsgSig.emit(self.taskMsg)

        except Exception as e:
            msg = f"{self.save_file.stem} - {ts_name}下载失败：{e}"
            self.taskMsg.error_msg += f"{msg}\n"
            logger.error(msg)
            # 重试 4 次
            if retries <= 3:
                await asyncio.sleep(1)
                await self.fetch(ts_url, ts_name, key, client, retries=retries + 1)

    async def run_download(self):
        self.start_time = datetime.now()
        async with asyncio.Semaphore(4):
            async with AsyncClient(
                    headers=self.headers,
                    timeout=self.timeout
            ) as client:
                tasks = [
                    self.fetch(url, f'{str(index).zfill(6)}.ts', self.decrypt_key, client)
                    for index, url in enumerate(self.ts_urls)
                ]
                await asyncio.gather(*tasks)

    def run(self):
        if self.isStop:
            self.taskMsg.status = TaskStatus.CANCEL
            self.taskMsg.message = "任务已取消"
            signalBus.taskMsgSig.emit(self.taskMsg)
            return

        if self.save_file.exists():
            self.taskMsg.message = '文件已存在，跳过下载'
            self.taskMsg.status = TaskStatus.CANCEL
            signalBus.taskMsgSig.emit(self.taskMsg)
            return
        else:
            self.temp_dir.mkdir(exist_ok=True, parents=True)
            self.ffmpeg_file.touch(exist_ok=True)

        # 获取最高画质的url
        self.taskMsg.status = TaskStatus.DOING
        self.taskMsg.message = '正在解析m3u8文件'
        signalBus.taskMsgSig.emit(self.taskMsg)
        try:
            m3u8_url = redirected_resolution_url(self.url, headers=self.headers, timeout=self.timeout)
            self.ts_urls, self.decrypt_key = m3u8_to_absolute_uris(m3u8_url, headers=self.headers, timeout=self.timeout)
            self.total_num = len(self.ts_urls)
        except Exception as e:
            self.taskMsg.error_msg += f"m3u8解析失败：{e}\n"
            self.taskMsg.status = TaskStatus.ERROR
            self.taskMsg.message = f'm3u8解析失败, {e}'
            signalBus.taskMsgSig.emit(self.taskMsg)
            logger.error(f"m3u8解析失败：{e}")
            return
        # 创建异步事件循环
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        # 下载ts文件
        try:
            self.__mainTask = self.loop.create_task(self.run_download())
            self.loop.run_until_complete(self.__mainTask)
        except asyncio.CancelledError as e:
            msg = f"任务取消：{self.save_file.stem}, {e}"
            self.taskMsg.error_msg += f"{msg}\n"
            self.taskMsg.status = TaskStatus.PAUSE
            signalBus.taskMsgSig.emit(self.taskMsg)
        finally:
            self.loop.run_until_complete(self.loop.shutdown_asyncgens())
            self.loop.close()
        del self.loop
        self.loop = None
        if self.isStop:
            return

        total = self.total_num
        if not total - 10 <= self.download_num <= total:
            msg = f"下载失败，已下载 {self.download_num} 个数据块，总数据块数为 {total} 个，ts文件损坏过多，请重试或自行合并"
            self.taskMsg.error_msg += msg
            self.taskMsg.message = msg
            self.taskMsg.status = TaskStatus.ERROR

            logger.error(msg)
            signalBus.taskMsgSig.emit(self.taskMsg)
            return
        # 合并ts文件
        self.taskMsg.message = '正在合并ts文件'
        signalBus.taskMsgSig.emit(self.taskMsg)
        is_success, error_string = merge_ts_to_mp4(self.ffmpeg_file, self.save_file, ffmpeg='ffmpeg')
        if not is_success:
            self.taskMsg.error_msg += error_string
            self.taskMsg.status = TaskStatus.ERROR
            logger.error(error_string)
            signalBus.taskMsgSig.emit(self.taskMsg)
            return
        self.taskMsg.message = '合并完成'
        signalBus.taskMsgSig.emit(self.taskMsg)
        # 删除临时文件夹
        shutil.rmtree(self.temp_dir)
        cover = get_video_cover(self.save_file)
        cover_path = (COVER_DIR / f'{self.tid}.png').as_posix()
        cover.save(cover_path)
        self.taskMsg.cover = cover_path
        self.taskMsg.status = TaskStatus.DONE
        signalBus.taskMsgSig.emit(self.taskMsg)

    def stop(self):
        self.isStop = True
        if self.__mainTask:
            self.__mainTask.cancel()
        self.taskMsg.message = '任务已暂停'
        self.taskMsg.status = TaskStatus.PAUSE
        signalBus.taskMsgSig.emit(self.taskMsg)

    def start(self):
        runner = M3u8Runner(
            self.task
        )
        signalBus.taskReAddSig.emit(runner)
        return runner
